#include <stdio.h>
#include <string.h>
// #include <time.h>
#include <sys/time.h>
#include "YoloV5.h"
#include "rknn_api.h"
#include "postprocess.h"

double __get_us(struct timeval t) { return (t.tv_sec * 1000000 + t.tv_usec); }

static void printRKNNTensor(rknn_tensor_attr *attr)
{
    printf("index=%d name=%s n_dims=%d dims=[%d %d %d %d] n_elems=%d size=%d "
           "fmt=%d type=%d qnt_type=%d fl=%d zp=%d scale=%f\n",
           attr->index, attr->name, attr->n_dims, attr->dims[3], attr->dims[2],
           attr->dims[1], attr->dims[0], attr->n_elems, attr->size, 0, attr->type,
           attr->qnt_type, attr->fl, attr->zp, attr->scale);
}

YoloV5::YoloV5()
{
}

YoloV5::~YoloV5()
{
    if(m_rknnCtx != 0){
        rknn_destroy(m_rknnCtx);
        m_rknnCtx = 0;
    }
}

int YoloV5::open(int modelWidth, int modelHeight, int clsNum)
{
    m_requestImgWidth = modelWidth;
    m_requestImgHeight = modelHeight;
    m_clssNum = clsNum;
}

int YoloV5::init(const unsigned char* modelData, int modelSize){
    if(!modelData || modelSize <= 0) return -1;
    int ret;
    /* Create the neural network */
    ret = rknn_init((rknn_context*)&m_rknnCtx, (void*)modelData, modelSize, 0, NULL);
    if (ret < 0){
        printf("rknn_init error ret=%d\n", ret);
        return -1;
    }

    rknn_sdk_version version;
    ret = rknn_query(m_rknnCtx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version));
    if (ret < 0){
        printf("rknn_init error ret=%d\n", ret);
        return -1;
    }

    printf("sdk version: %s driver version: %s\n", version.api_version, version.drv_version);

    rknn_input_output_num io_num;
    ret = rknn_query(m_rknnCtx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret < 0){
        printf("rknn_init error ret=%d\n", ret);
        return -1;
    }

    printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);

    m_tensorInputNum = io_num.n_input;
    m_tensorOutputNum = io_num.n_output;

    rknn_tensor_attr input_attrs[io_num.n_input];
    memset(input_attrs, 0, sizeof(input_attrs));
    for (int i = 0; i < io_num.n_input; i++){
        input_attrs[i].index = i;
        ret = rknn_query(m_rknnCtx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret < 0){
            printf("rknn_init error ret=%d\n", ret);
            return -1;
        }
        printRKNNTensor(&(input_attrs[i]));
    }

    rknn_tensor_attr output_attrs[m_tensorOutputNum];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (int i = 0; i < m_tensorOutputNum; i++){
        output_attrs[i].index = i;
        ret = rknn_query(m_rknnCtx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]),
                         sizeof(rknn_tensor_attr));
        printRKNNTensor(&(output_attrs[i]));
    }

    if (input_attrs[0].fmt == RKNN_TENSOR_NCHW){
        printf("model is NCHW input fmt\n");
        m_requestImgWidth = input_attrs[0].dims[0];
        m_requestImgHeight = input_attrs[0].dims[1];
    }
    else{
        printf("model is NHWC input fmt\n");
        m_requestImgWidth = input_attrs[0].dims[1];
        m_requestImgHeight = input_attrs[0].dims[2];
    }

    printf("model input height=%d, width=%d, channel=%d\n", m_requestImgHeight, m_requestImgWidth,
           m_requestImgChannel);

    return 0;
}

int YoloV5::detect(const unsigned char* pData, std::vector<BBOX>& detectResults){
    if(m_rknnCtx == 0 || !pData)   return -1;

    int ret = 0;
    struct timeval start_time, stop_time;
    const float nms_threshold = 0.45;
    const float conf_threshold = 0.25;
    
    rknn_tensor_attr output_attrs[m_tensorOutputNum];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (int i = 0; i < m_tensorOutputNum; i++)
    {
        output_attrs[i].index = i;
        ret = rknn_query(m_rknnCtx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]),
                         sizeof(rknn_tensor_attr));
        // printRKNNTensor(&(output_attrs[i]));
    }
    //
    rknn_input inputs[1];
    memset(inputs, 0, sizeof(inputs));
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = m_requestImgWidth * m_requestImgHeight * m_requestImgChannel;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].pass_through = 0;
    inputs[0].buf = (void*)pData;

    // gettimeofday(&start_time, NULL);
    rknn_inputs_set(m_rknnCtx, m_tensorInputNum, inputs);

    rknn_output outputs[m_tensorOutputNum];
    memset(outputs, 0, sizeof(outputs));
    for (int i = 0; i < m_tensorOutputNum; i++)
    {
        outputs[i].index = i;
        outputs[i].want_float = 0;
    }

    ret = rknn_run(m_rknnCtx, NULL);
    gettimeofday(&stop_time, NULL);
    // printf("once run use %f ms\n",
    //        (__get_us(stop_time) - __get_us(start_time)) / 1000);

    ret = rknn_outputs_get(m_rknnCtx, m_tensorOutputNum, outputs, NULL);
    // gettimeofday(&stop_time, NULL);

    //post process
    float scale_w = (float)m_requestImgWidth / m_requestImgWidth;
    float scale_h = (float)m_requestImgHeight / m_requestImgHeight;

    detect_result_group_t detect_result_group;
    std::vector<float> out_scales;
    std::vector<int32_t> out_zps;
    for (int i = 0; i < m_tensorOutputNum; ++i)
    {
        out_scales.push_back(output_attrs[i].scale);
        out_zps.push_back(output_attrs[i].zp);
    }
    
    BOX_RECT pads;
    memset(&pads, 0, sizeof(BOX_RECT));
    post_process((int8_t *)outputs[0].buf, (int8_t *)outputs[1].buf, (int8_t *)outputs[2].buf, m_requestImgHeight, m_requestImgWidth,
            conf_threshold, nms_threshold, pads, scale_w, scale_h, out_zps, out_scales, &detect_result_group, m_clssNum);
    for (int i = 0; i < detect_result_group.count; i++)
    {
        detect_result_t *det_result = &(detect_result_group.results[i]);
        // sprintf(text, "%s %.2f", det_result->name, det_result->prop);
        if(1){
            printf("id:%d @ (%d %d %d %d)\n",
                det_result->id,
                det_result->box.left, det_result->box.top, det_result->box.right, det_result->box.bottom);
        }

        BBOX box{};
        box.rect.ltX = det_result->box.left;
        box.rect.ltY = det_result->box.top;
        box.rect.rbX = det_result->box.right;
        box.rect.rbY = det_result->box.bottom;
        box.id = det_result->id;
        box.score = det_result->prop * 100;

        detectResults.push_back(box);
    }
    
    ret = rknn_outputs_release(m_rknnCtx, m_tensorOutputNum, outputs);
    if(ret != 0){
        printf("rknn_outputs_release failed! \n");
    }

    return 0;
}

void YoloV5::getModelRequstSize(int& width, int& height){
    width = m_requestImgWidth;
    height = m_requestImgHeight;
}
